import pandas as pd
import numpy as np
from pathlib import Path
import logging
from datetime import datetime
from sklearn.preprocessing import LabelEncoder

'''
特征工程分析
1. 编码转换
    Medal_encoding.csv: 将奖牌类型转为数值编码
    NOC_encoding.csv: 将国家/地区代码转为数值编码
    Sport_encoding.csv: 将运动项目转为数值编码
    这些编码转换的目的是将分类变量转换为机器学习模型可以处理的数值形式。

2. 特征提取
    athlete_features.csv: 包含运动员个人特征
        参赛经验(participation_count)
        历史成绩(avg_performance, best_performance)
        专项程度(sport_count, event_count)
    country_features.csv: 包含国家层面特征
        历史奖牌统计
        参赛历史
        主办经验
    sport_features.csv: 包含项目特征
        项目历史数据
        参与国家数
        竞争指数
    time_series_features.csv: 包含时间序列特征
        奖牌趋势
        年度统计
'''

class OlympicsFeatureEngineering:
    def __init__(self):
        """初始化数据加载"""
        self.athletes_df = pd.read_csv('summerOly_athletes.csv')
        self.medals_df = pd.read_csv('summerOly_medal_counts.csv')
        self.hosts_df = pd.read_csv('summerOly_hosts.csv')
        self.programs_df = pd.read_csv('summerOly_programs.csv')

        # 创建特征数据保存目录
        self.features_dir = Path('features')
        self.features_dir.mkdir(exist_ok=True)

        # 初始化标签编码器
        self.label_encoders = {}

    def create_athlete_features(self):
        """构建运动员相关特征"""
        athlete_features = self.athletes_df.copy()
        athlete_features = athlete_features.rename(columns={'Team': 'country'})  # 重命名Team列为country

        # 1. 运动员参赛经验特征 - 按运动员和国家分组
        experience = athlete_features.groupby(['Name', 'country']).agg({
            'Year': ['count', 'min', 'max']
        }).reset_index()
        experience.columns = ['Name', 'country', 'participation_count', 'first_year', 'last_year']

        # 2. 运动员历史成绩特征
        medal_mapping = {'Gold': 3, 'Silver': 2, 'Bronze': 1, 'No medal': 0}
        athlete_features['medal_value'] = athlete_features['Medal'].map(medal_mapping)

        performance = athlete_features.groupby(['Name', 'country']).agg({
            'medal_value': ['mean', 'max', 'sum']
        }).reset_index()
        performance.columns = ['Name', 'country', 'avg_performance', 'best_performance', 'total_medals']

        # 3. 运动员专项特征
        specialization = athlete_features.groupby(['Name', 'country']).agg({
            'Sport': 'nunique',
            'Event': 'nunique'
        }).reset_index()
        specialization.columns = ['Name', 'country', 'sport_count', 'event_count']

        # 合并所有运动员特征
        athlete_features = pd.merge(experience, performance, on=['Name', 'country'], how='left')
        athlete_features = pd.merge(athlete_features, specialization, on=['Name', 'country'], how='left')

        # 确保列的顺序，将 Name 和 country 放在最前面
        columns_order = ['Name', 'country'] + [col for col in athlete_features.columns 
                                             if col not in ['Name', 'country']]
        athlete_features = athlete_features[columns_order]

        # 保存特征
        athlete_features.to_csv(self.features_dir / 'athlete_features.csv', index=False)
        return athlete_features

    def create_country_features(self):
        """构建国家相关特征"""
        # 1. 历史奖牌统计特征
        country_features = self.medals_df.groupby('country').agg({
            'Gold': ['sum', 'mean'],
            'Silver': ['sum', 'mean'],
            'Bronze': ['sum', 'mean'],
            'Total': ['sum', 'mean'],
            'Year': ['count', 'min', 'max']
        }).reset_index()

        # 重命名列
        column_names = ['country', 'total_gold', 'avg_gold', 'total_silver', 'avg_silver',
                        'total_bronze', 'avg_bronze', 'total_medals', 'avg_medals',
                        'participation_years', 'first_year', 'last_year']
        country_features.columns = column_names

        # 2. 项目优势特征
        sport_strength = self.athletes_df[self.athletes_df['Medal'] != 'No medal']
        sport_strength = sport_strength.groupby(['country', 'Sport']).size().reset_index(name='medals_in_sport')

        # 获取每个国家最强的前3个项目
        top_sports = sport_strength.sort_values('medals_in_sport', ascending=False) \
            .groupby('country').head(3)

        # 3. 主办国经验
        host_experience = self.hosts_df[self.hosts_df['Host'] != 'Cancelled']
        host_experience = host_experience.groupby('Country').size().reset_index(name='times_as_host')

        # 合并特征
        country_features = pd.merge(country_features,
                                    host_experience,
                                    left_on='country',
                                    right_on='Country',
                                    how='left')
        country_features['times_as_host'] = country_features['times_as_host'].fillna(0)

        # 保存特征
        country_features.to_csv(self.features_dir / 'country_features.csv', index=False)
        return country_features

    def create_sport_features(self):
        """构建运动项目相关特征"""
        # 1. 项目历史特征
        sport_features = self.athletes_df.groupby('Sport').agg({
            'Year': ['nunique', 'min', 'max'],  # 举办次数和时间跨度
            'NOC': 'nunique',  # 参与国家数
            'Name': 'count'  # 参与人数
        }).reset_index()

        # 重命名列
        sport_features.columns = ['Sport', 'years_held', 'first_year',
                                  'last_year', 'participating_countries',
                                  'total_participants']

        # 2. 项目竞争度特征
        medals_by_sport = self.athletes_df[self.athletes_df['Medal'] != 'No medal']
        medals_by_sport = medals_by_sport.groupby('Sport').agg({
            'NOC': 'nunique',  # 获奖国家数
            'Medal': 'count'  # 奖牌总数
        }).reset_index()

        medals_by_sport.columns = ['Sport', 'medal_winning_countries', 'total_medals']

        # 计算竞争度指标
        medals_by_sport['competition_index'] = medals_by_sport['medal_winning_countries'] / \
                                               medals_by_sport['total_medals']

        # 合并特征
        sport_features = pd.merge(sport_features, medals_by_sport, on='Sport', how='left')

        # 保存特征
        sport_features.to_csv(self.features_dir / 'sport_features.csv', index=False)
        return sport_features

    def create_time_series_features(self):
        """构建时间序列特征"""
        # 1. 按年份的全局趋势特征
        yearly_features = self.medals_df.groupby('Year').agg({
            'Total': 'sum',  # 总奖牌数
            'NOC': 'nunique'  # 参与国家数
        }).reset_index()

        # 2. 计算移动平均和趋势
        yearly_features['medals_ma_4'] = yearly_features['Total'].rolling(window=4).mean()
        yearly_features['medals_trend'] = yearly_features['Total'].pct_change(4)

        # 3. 添加奥运会举办地信息
        yearly_features = pd.merge(yearly_features,
                                   self.hosts_df[['Year', 'Country']],
                                   on='Year',
                                   how='left')

        # 保存特征
        yearly_features.to_csv(self.features_dir / 'time_series_features.csv', index=False)
        return yearly_features

    def encode_categorical_features(self):
        """对分类特征进行编码"""
        categorical_features = {
            'NOC': self.athletes_df['NOC'].unique(),
            'Sport': self.athletes_df['Sport'].unique(),
            'Medal': ['No medal', 'Bronze', 'Silver', 'Gold']
        }

        # 为每个分类特征创建编码器
        for feature, values in categorical_features.items():
            le = LabelEncoder()
            le.fit(values)
            self.label_encoders[feature] = le

        # 保存编码映射
        encoding_dict = {
            feature: dict(zip(le.classes_, le.transform(le.classes_)))
            for feature, le in self.label_encoders.items()
        }

        # 将编码映射保存为CSV
        for feature, mapping in encoding_dict.items():
            pd.DataFrame(list(mapping.items()),
                         columns=[feature, f'{feature}_encoded']) \
                .to_csv(self.features_dir / f'{feature}_encoding.csv', index=False)

    def run_feature_engineering(self):
        """运行所有特征工程流程"""
        print("开始特征工程...")

        print("1. 构建运动员特征...")
        athlete_features = self.create_athlete_features()

        print("2. 构建国家特征...")
        country_features = self.create_country_features()

        print("3. 构建运动项目特征...")
        sport_features = self.create_sport_features()

        print("4. 构建时间序列特征...")
        time_series_features = self.create_time_series_features()

        print("5. 编码分类特征...")
        self.encode_categorical_features()

        print("特征工程完成！所有特征文件已保存到 features 目录")

        return {
            'athlete_features': athlete_features,
            'country_features': country_features,
            'sport_features': sport_features,
            'time_series_features': time_series_features
        }


def main():
    feature_engineering = OlympicsFeatureEngineering()
    features = feature_engineering.run_feature_engineering()

    # 打印特征统计信息
    for feature_name, feature_df in features.items():
        print(f"\n{feature_name} shape:", feature_df.shape)
        print("\nSample features:")
        print(feature_df.head())


if __name__ == "__main__":
    main()
